#include<opencv2/opencv.hpp>
#include <iostream>
#include<random>
#include<unordered_set>
#include<time.h>

using namespace std;

cv::Mat rls(vector<double> t, vector<double> data) {
    int order = 2;
    cv::Mat A = cv::Mat::zeros(t.size(), order, CV_64FC1);
    for (int row = 0; row < A.rows; row++) {
        for (int col = 0; col < A.cols; col++) {
            A.at<double>(row, col) = pow(t[row], col);
        }
    }
    // cout<<"A "<<A<<endl;
    // 构建B矩阵
    cv::Mat B = cv::Mat::zeros(t.size(), 1, CV_64FC1);
    for (int row = 0; row < B.rows; row++) {
        B.at<double>(row, 0)= data[row];
    }
    cv::Mat X;
    solve(A, B, X, CV_SVD);
    return X;
}

auto  motion_model = [](double t, cv::Mat beta ) {
    return beta.at<double>(1, 0)*t+beta.at<double>(0, 0);
};

void ransac(vector<double> t, vector<double> data,cv::Mat& model_ret){
    int sample_selcect_size_ = 10;
    int iter_max_times_= 50;
    int inliner_size_threshold_ = 10;
    double threshold_ = 10;
    cv::Mat best_model_;
    double best_error_ = 1e10;
    default_random_engine  gen_(10);
    uniform_int_distribution<int> dist_(0,data.size()-1);
    unordered_set<int> us_;

    bool is_converge_ = false;

    auto error_fun = [](double y_pre,double y){
        return pow(y_pre-y,2);
    };

    while(iter_max_times_){
        iter_max_times_--;
        vector<double> t_se,data_se,t_less,data_less;
        while(us_.size()<sample_selcect_size_){
            int idx = dist_(gen_);
            if(us_.count(idx)==0){
                us_.insert(idx);
            }
        }

        cout<<"time "<<iter_max_times_<<endl;
        for(int i=0;i<data.size();i++){
            if(us_.count(i)){
                // cout<<i<<endl;
                t_se.emplace_back(t[i]);
                data_se.emplace_back(data[i]);
            }else{
                t_less.emplace_back(t[i]);
                data_less.emplace_back(data[i]);
            }
        }

        us_.clear();

        cv::Mat X = rls(t_se,data_se);

        vector<int> inliner_idx;
        for(int i=0;i<t_less.size();i++){ 
            double error = error_fun(motion_model(t_less[i],X),data_less[i]);
            if(error<threshold_){
                inliner_idx.emplace_back(i);
            }
        }

        if(inliner_idx.size()>inliner_size_threshold_){
            for(auto& idx:inliner_idx){
                t_se.emplace_back(t[idx]);
                data_se.emplace_back(data[idx]);
            }

            X = rls(t_se,data_se);

            double mean_error = 0;
            for(int i=0;i<t_se.size();i++){
                mean_error = mean_error + error_fun(motion_model(t_se[i],X),data_se[i]);
            }
            mean_error = mean_error/t_se.size();

            if(mean_error<best_error_){
                is_converge_ = true;
                best_error_ = mean_error;
                best_model_ = X;
                cout<<"best error "<<to_string(best_error_)<<endl;
            }
        }
    }
    model_ret = best_model_;
    cout<<"best model "<<best_model_<<endl;
}

int main(){
    default_random_engine gen;
    normal_distribution<double> normal_dist(0,10);

    vector<double> t;
    vector<double> data;
    vector<double> data_with_noise;

    cv::Mat X_target(2,1,CV_64F);
    X_target.at<double>(1,0) = 2;
    X_target.at<double>(0,0) = 1;

    for(int i=0;i<400;i++){
        t.emplace_back(i);
        data.emplace_back(motion_model(i,X_target));
        data_with_noise.emplace_back(motion_model(i,X_target)+normal_dist(gen));
    }

    cv::Mat show_image(800,400,CV_8UC3,cv::Scalar(0,0,0));
    for(int i=0;i<400;i++){
        circle(show_image, cv::Point2d(i,data_with_noise[i]), 3, cv::Scalar(255,0,0), 2);
    }

    cv::Mat X_rls(2,1,CV_64F);
    X_rls.at<double>(1,0) = 0;
    X_rls.at<double>(0,0) = 0;
    ransac(t,data_with_noise,X_rls);



    cv::line(show_image,cv::Point2d(1,motion_model(1,X_rls)),cv::Point2d(390,motion_model(390,X_rls)),cv::Scalar(0,0,255),2);


    cout<<X_rls<<endl;

    cv::imshow("test",show_image);
	cv::waitKey(0);
	return 0;
}
